import tensorflow as tf

# 加载图片，对图像进行归一化并按照指定大小进行缩放
def load_img(path):
    max_dim = 1080
    # 加载图像
    img = tf.io.read_file(path)
    # 图像解码
    img = tf.image.decode_image(img, channels=3)
    # 图像归一化
    img = tf.image.convert_image_dtype(img, tf.float32)
    # 获取图像形状
    shape = tf.cast(tf.shape(img)[:-1], tf.float32)
    long_dim = max(shape)
    scale = max_dim / long_dim
    # 根据比例缩放
    shape = tf.cast(shape * scale, tf.int32)    
    # 重载图像
    img = tf.image.resize(img, shape)
    img = img[tf.newaxis, :]
    return img

# 计算gram矩阵，使用爱因斯坦求和约束
def gram_matrix(input_tensor):
    result = tf.linalg.einsum('bijc,bijd->bcd', input_tensor, input_tensor)
    input_shape = tf.shape(input_tensor)
    num_locations = tf.cast(input_shape[1]*input_shape[2], tf.float32)
    return result/(num_locations)

# 精度截断
def clip_0_1(image):
    return tf.clip_by_value(image, clip_value_min=0.0, clip_value_max=1.0)

# 模糊操作，降低图像的像素锐化
def high_pass(image):
    x = image[:,:,1:,:] - image[:,:,:-1,:]
    y = image[:,1:,:,:] - image[:,:-1,:,:]
    return x, y
def total_variation_loss(image):
    x, y = high_pass(image)
    return tf.reduce_mean(x**2) + tf.reduce_mean(y**2)